# -*- coding: utf-8 -*-
"""
主程序
运行检索方法对比评估
"""

import os
import sys
import argparse
from typing import List, Dict, Any
from loguru import logger
import json
from datetime import datetime

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from config import (
    OLLAMA_CONFIG, RETRIEVAL_METHODS, EVALUATION_METRICS,
    TEST_DATASET_CONFIG, OUTPUT_CONFIG, LOGGING_CONFIG
)
from data_manager import DataManager
from evaluator import Evaluator
from ollama_client import ollama_client


def setup_logging():
    """设置日志"""
    # 移除默认处理器
    logger.remove()
    
    # 添加控制台处理器
    logger.add(
        sys.stdout,
        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
        level=LOGGING_CONFIG.get('level', 'INFO')
    )
    
    # 添加文件处理器
    log_file = LOGGING_CONFIG.get('file', 'logs/comparison.log')
    os.makedirs(os.path.dirname(log_file), exist_ok=True)
    
    logger.add(
        log_file,
        format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}",
        level=LOGGING_CONFIG.get('level', 'INFO'),
        rotation="10 MB",
        retention="7 days"
    )


def check_ollama_connection() -> bool:
    """检查Ollama连接"""
    logger.info("🔗 检查Ollama连接...")
    
    try:
        if ollama_client.check_connection():
            logger.info("✅ Ollama连接正常")
            return True
        else:
            logger.error("❌ 无法连接到Ollama服务")
            logger.info("💡 请确保:")
            logger.info("   1. Ollama服务正在运行")
            logger.info("   2. bge-m3:latest模型已下载 (用于embedding)")
            logger.info("   3. bge-rerank-v2-m3模型已下载 (用于rerank)")
            logger.info("   4. 服务地址正确")
            return False
    except Exception as e:
        logger.error(f"❌ Ollama连接异常: {e}")
        return False


def load_test_data(data_manager: DataManager, config: Dict[str, Any]) -> List[Dict[str, Any]]:
    """加载测试数据"""
    logger.info("📊 加载测试数据...")
    
    try:
        # 从配置获取数据源
        data_source = config.get('source', 'database')
        
        if data_source == 'file':
            # 从文件加载
            file_path = config.get('file_path', 'data/test_questions.json')
            test_data = data_manager.load_from_file(file_path)
        else:
            # 从数据库生成
            db_path = config.get('db_path', '../data/qa_database.db')
            sample_size = config.get('sample_size', 100)
            test_data = data_manager.load_from_database(db_path, sample_size)
        
        logger.info(f"✅ 加载测试数据完成，数据量: {len(test_data)}")
        return test_data
        
    except Exception as e:
        logger.error(f"❌ 加载测试数据失败: {e}")
        return []


def prepare_evaluation_data(test_data: List[Dict[str, Any]], 
                          train_ratio: float = 0.8) -> tuple:
    """准备评估数据"""
    logger.info("🔄 准备评估数据...")
    
    # 按比例分割训练和测试数据
    split_index = int(len(test_data) * train_ratio)
    
    train_data = test_data[:split_index]
    test_questions = test_data[split_index:]
    
    logger.info(f"📊 数据分割完成 - 训练: {len(train_data)}, 测试: {len(test_questions)}")
    
    return train_data, test_questions


def run_evaluation(methods: List[str] = None, 
                  sample_size: int = None,
                  output_dir: str = None) -> Dict[str, Any]:
    """运行评估"""
    logger.info("🚀 开始运行检索方法对比评估")
    
    # 1. 检查Ollama连接
    if not check_ollama_connection():
        logger.warning("⚠️ Ollama连接失败，某些检索方法可能无法正常工作")
        logger.info("继续运行其他检索方法的评估...")
    
    # 2. 初始化数据管理器
    data_manager = DataManager()
    
    # 3. 加载测试数据
    test_data = load_test_data(data_manager, TEST_DATASET_CONFIG)
    if not test_data:
        logger.error("❌ 无法加载测试数据")
        return {}
    
    # 限制样本大小
    if sample_size and sample_size < len(test_data):
        test_data = test_data[:sample_size]
        logger.info(f"📊 限制样本大小为: {sample_size}")
    
    # 4. 准备评估数据
    train_data, test_questions = prepare_evaluation_data(test_data)
    
    # 5. 初始化评估器
    evaluator_config = {
        'output_dir': output_dir or OUTPUT_CONFIG.get('results_dir', 'results'),
        'enable_detailed_analysis': True,
        'enable_visualization': True
    }
    evaluator = Evaluator(evaluator_config)
    
    # 6. 初始化检索器
    logger.info("🔧 初始化检索器...")
    if not evaluator.initialize_retrievers(train_data):
        logger.error("❌ 检索器初始化失败")
        return {}
    
    # 7. 运行评估
    logger.info("📊 开始评估...")
    evaluation_results = evaluator.evaluate_all_methods(test_questions, methods)
    
    if not evaluation_results:
        logger.error("❌ 评估失败")
        return {}
    
    # 8. 生成报告
    logger.info("📄 生成评估报告...")
    report_content = evaluator.generate_report()
    
    # 9. 保存详细结果
    logger.info("💾 保存详细结果...")
    detailed_file = evaluator.save_detailed_results()
    
    # 10. 创建可视化
    logger.info("📊 创建可视化图表...")
    chart_files = evaluator.create_visualizations()
    
    # 11. 汇总结果
    summary = {
        'evaluation_time': datetime.now().isoformat(),
        'methods_evaluated': list(evaluation_results.keys()),
        'test_data_size': len(test_questions),
        'train_data_size': len(train_data),
        'report_content': report_content,
        'detailed_results_file': detailed_file,
        'chart_files': chart_files,
        'metrics_summary': {}
    }
    
    # 添加指标摘要
    for method_name, result in evaluation_results.items():
        summary['metrics_summary'][method_name] = {
            'precision': result.metrics.get('precision', 0),
            'recall': result.metrics.get('recall', 0),
            'f1': result.metrics.get('f1', 0),
            'avg_response_time': result.metrics.get('avg_response_time', 0),
            'success_rate': result.metrics.get('success_rate', 0)
        }
    
    logger.info("✅ 评估完成")
    return summary


def print_summary(summary: Dict[str, Any]):
    """打印评估摘要"""
    if not summary:
        return
    
    print("\n" + "="*80)
    print("🎯 检索方法对比评估结果摘要")
    print("="*80)
    
    print(f"\n📅 评估时间: {summary['evaluation_time']}")
    print(f"📊 评估方法: {', '.join(summary['methods_evaluated'])}")
    print(f"📝 测试数据量: {summary['test_data_size']}")
    print(f"🔧 训练数据量: {summary['train_data_size']}")
    
    print("\n📈 性能指标摘要:")
    print("-" * 80)
    print(f"{'方法':<15} {'精度':<8} {'召回率':<8} {'F1分数':<8} {'响应时间(s)':<12} {'成功率':<8}")
    print("-" * 80)
    
    for method, metrics in summary['metrics_summary'].items():
        print(f"{method:<15} {metrics['precision']:<8.3f} {metrics['recall']:<8.3f} "
              f"{metrics['f1']:<8.3f} {metrics['avg_response_time']:<12.3f} "
              f"{metrics['success_rate']:<8.3f}")
    
    print("\n📁 输出文件:")
    if summary.get('detailed_results_file'):
        print(f"  📄 详细结果: {summary['detailed_results_file']}")
    
    if summary.get('chart_files'):
        print("  📊 可视化图表:")
        for chart_file in summary['chart_files']:
            print(f"    - {chart_file}")
    
    print("\n" + "="*80)


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='检索方法对比评估')
    
    parser.add_argument(
        '--methods', 
        nargs='+', 
        choices=['vector', 'keyword', 'hybrid', 'knowledge_graph', 'rerank', 'multi_step', 'structured'],
        help='要评估的检索方法，默认评估所有方法'
    )
    
    parser.add_argument(
        '--sample-size', 
        type=int, 
        help='测试样本大小，默认使用所有数据'
    )
    
    parser.add_argument(
        '--output-dir', 
        type=str, 
        help='输出目录，默认为results'
    )
    
    parser.add_argument(
        '--log-level', 
        choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'],
        default='INFO',
        help='日志级别'
    )
    
    parser.add_argument(
        '--no-visualization', 
        action='store_true',
        help='禁用可视化图表生成'
    )
    
    args = parser.parse_args()
    
    # 设置日志级别
    LOGGING_CONFIG['level'] = args.log_level
    setup_logging()
    
    logger.info("🎯 启动检索方法对比评估程序 (纯检索对比，无LLM问答)")
    logger.info(f"📋 评估参数: {vars(args)}")
    logger.info(f"🤖 Embedding模型: {OLLAMA_CONFIG.get('embedding_model', 'bge-m3:latest')}")
    logger.info(f"🔄 Rerank模型: {OLLAMA_CONFIG.get('rerank_model', 'bge-rerank-v2-m3')}")
    
    try:
        # 运行评估
        summary = run_evaluation(
            methods=args.methods,
            sample_size=args.sample_size,
            output_dir=args.output_dir
        )
        
        # 打印摘要
        print_summary(summary)
        
        # 保存摘要
        if summary:
            output_dir = args.output_dir or OUTPUT_CONFIG.get('results_dir', 'results')
            summary_file = os.path.join(output_dir, f"evaluation_summary_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
            
            try:
                with open(summary_file, 'w', encoding='utf-8') as f:
                    json.dump(summary, f, ensure_ascii=False, indent=2)
                logger.info(f"💾 评估摘要已保存: {summary_file}")
            except Exception as e:
                logger.warning(f"⚠️ 保存评估摘要失败: {e}")
        
        logger.info("🎉 程序执行完成")
        
    except KeyboardInterrupt:
        logger.warning("⚠️ 用户中断程序执行")
        sys.exit(1)
    except Exception as e:
        logger.error(f"❌ 程序执行异常: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()